# -*- coding: utf-8 -*-
import os
import pandas as pd
import numpy as np
import sys
import efinance as ef
from pypinyin import pinyin, Style

from sklearn.preprocessing import RobustScaler
from tensorflow.keras.models import Sequential
from tensorflow.keras.layers import LSTM, Dense, Dropout
from tensorflow.keras.callbacks import EarlyStopping
import tensorflow as tf

import matplotlib.pyplot as plt

# --- 智能GPU配置 ---
def configure_gpu(force_cpu=False):
    """配置GPU设备，支持CPU回退"""
    if force_cpu:
        os.environ['CUDA_VISIBLE_DEVICES'] = ''
        print("ℹ️ 强制使用CPU模式进行训练")
        return False
        
    gpus = tf.config.experimental.list_physical_devices('GPU')
    if gpus:
        try:
            # 允许GPU内存增长
            for gpu in gpus:
                tf.config.experimental.set_memory_growth(gpu, True)
            print(f"✅ GPU配置成功！检测到 {len(gpus)} 个GPU设备")
            for i, gpu in enumerate(gpus):
                print(f"   GPU {i}: {gpu.name}")
            print("⚠️ 注意：RTX 5090首次使用可能需要30分钟JIT编译CUDA核心")
            return True
        except Exception as e:
            print(f"⚠️ GPU配置失败: {e}")
            print("   自动回退到CPU模式")
            os.environ['CUDA_VISIBLE_DEVICES'] = ''
            return False
    else:
        print("ℹ️ 未检测到GPU设备，使用CPU进行训练")
        return False

# 检查环境变量，允许用户选择CPU模式
force_cpu_mode = os.getenv('FORCE_CPU', 'false').lower() == 'true'
gpu_available = configure_gpu(force_cpu=force_cpu_mode)

def to_pinyin(text):
    s = ''
    for i in pinyin(text, style=Style.NORMAL):
        s += ''.join(i)
    return s

def create_dataset(X, y, time_steps=1):
    """ 将数据转换为时间序列样本 """
    Xs, ys = [], []
    for i in range(len(X) - time_steps):
        v = X.iloc[i:(i + time_steps)].values
        Xs.append(v)
        ys.append(y.iloc[i + time_steps])
    return np.array(Xs), np.array(ys)

# --- 1. 获取参数 ---
if len(sys.argv) < 3:
    print("用法: python3 04b_train_lstm.py <股票代码> <frequency>")
    sys.exit(1)
stock_code = sys.argv[1]
frequency = sys.argv[2]

# --- 2. 获取公司名称 ---
stock_name_pinyin = stock_code
try:
    quote_df = ef.stock.get_quote_history(stock_code, klt=101, end='20300101', limit=2)
    if not quote_df.empty and '股票名称' in quote_df.columns:
        stock_name_pinyin = to_pinyin(quote_df['股票名称'].iloc[0])
except Exception:
    pass

# --- 3. 定义输入文件名 ---
input_filename = f'output/data/stock_data_{stock_code}_{frequency}_with_indicators.csv'
model_filename = f'output/models/lstm_model_{stock_code}_{frequency}.h5'
plot_filename = f'output/plots/prediction_vs_actual_{stock_code}_{frequency}_lstm.png'

# --- 4. 加载数据 ---
try:
    print(f"开始加载: {input_filename}...")
    # 关键修复：现在文件使用英文列名，所以用 'date' 作为索引列
    df = pd.read_csv(input_filename, index_col='date', parse_dates=True)
    print("加载成功。")
except FileNotFoundError:
    print(f"错误: 找不到数据文件 {input_filename}。")
    sys.exit(1)
except ValueError:
    print(f"错误: {input_filename} 中似乎找不到 'date' 列。请确保文件格式正确。")
    sys.exit(1)

# --- 5. 准备特征和目标变量 ---
df.dropna(inplace=True)
df.replace([np.inf, -np.inf], np.nan, inplace=True)
df.dropna(inplace=True)

# 关键修复：从特征中排除'open'和'close'，因为它们是预测目标，以避免数据泄漏
features = [col for col in df.columns if col not in ['date', 'code', 'target', 'open', 'close']]
X_df = df[features]
# 升级为双目标预测：同时预测开盘价和收盘价
y_df = df[['open', 'close']]

# --- 6. 数据归一化 ---
# 归一化输入特征
X_scaler = RobustScaler()
X_scaled = X_scaler.fit_transform(X_df)

# 关键修复：归一化目标变量（开盘价和收盘价）
y_scaler = RobustScaler()
y_scaled = y_scaler.fit_transform(y_df)

# 将归一化后的数据转回DataFrame以保留列名
X_scaled_df = pd.DataFrame(X_scaled, index=X_df.index, columns=features)
y_scaled_df = pd.DataFrame(y_scaled, index=y_df.index, columns=['open', 'close'])

# --- 7. 创建时间序列数据集 ---
TIME_STEPS = 60
# 修复：使用归一化后的目标变量
X_ts, y_ts = create_dataset(X_scaled_df, y_scaled_df, TIME_STEPS)

# --- 8. 划分训练集和测试集 ---
train_size = int(len(X_ts) * 0.8)
X_train, X_test = X_ts[:train_size], X_ts[train_size:]
y_train, y_test = y_ts[:train_size], y_ts[train_size:]

# --- 9. 构建和训练LSTM模型 ---
# 从训练数据中获取时间步数和特征数
time_steps = X_train.shape[1]
n_features = X_train.shape[2]

model = Sequential([
    LSTM(50, return_sequences=False, input_shape=(time_steps, n_features)),
    Dropout(0.2),
    # 升级：输出层改为2个神经元，分别对应 open 和 close 的预测
    Dense(2)
])

model.compile(optimizer='adam', loss='mean_squared_error')
model.summary()

early_stopping = EarlyStopping(monitor='val_loss', patience=10, restore_best_weights=True)

history = model.fit(X_train, y_train,
                    epochs=100,
                    batch_size=32,
                    validation_split=0.2,
                    callbacks=[early_stopping],
                    verbose=1)

# --- 10. 评估模型 ---
test_loss = model.evaluate(X_test, y_test, verbose=0)
print(f"模型的均方根误差 (RMSE) 是: {np.sqrt(test_loss):.4f}")

# --- 11. 保存模型和Scaler ---
# 确保output目录存在
os.makedirs('output/models', exist_ok=True)
os.makedirs('output/scalers', exist_ok=True)

model.save(model_filename)
print(f"LSTM模型已成功保存到: {model_filename}")

# 保存归一化器，用于预测时的数据处理
import joblib
X_scaler_filename = f'output/scalers/X_scaler_{stock_code}_{frequency}.pkl'
y_scaler_filename = f'output/scalers/y_scaler_{stock_code}_{frequency}.pkl'
joblib.dump(X_scaler, X_scaler_filename)
joblib.dump(y_scaler, y_scaler_filename)
print(f"特征归一化器已保存到: {X_scaler_filename}")
print(f"目标归一化器已保存到: {y_scaler_filename}")

# 可视化预测结果
predictions_scaled = model.predict(X_test)
# 关键修复：将预测结果反归一化回原始价格尺度
predictions = y_scaler.inverse_transform(predictions_scaled)
# 同样将真实值反归一化
y_test_original = y_scaler.inverse_transform(y_test)

# 提取开盘价和收盘价进行可视化（使用反归一化后的真实价格）
y_test_open = y_test_original[:, 0]
y_test_close = y_test_original[:, 1]
predictions_open = predictions[:, 0]
predictions_close = predictions[:, 1]

# 创建图表
plt.figure(figsize=(14, 7))
# 绘制收盘价
plt.plot(df.index[len(df) - len(y_test_close):], y_test_close, color='red', label='Actual Close Price')
plt.plot(df.index[len(df) - len(y_test_close):], predictions_close, color='orange', linestyle='--', label='Predicted Close Price')
# 绘制开盘价
plt.plot(df.index[len(df) - len(y_test_open):], y_test_open, color='blue', label='Actual Open Price')
plt.plot(df.index[len(df) - len(y_test_open):], predictions_open, color='cyan', linestyle='--', label='Predicted Open Price')

plt.title(f'{stock_name_pinyin} ({stock_code}) - Open & Close Price Prediction (LSTM)')
plt.xlabel('Date')
plt.ylabel('Stock Price')
plt.legend()
plt.grid(True)
# 确保output/plots目录存在
os.makedirs('output/plots', exist_ok=True)
plt.savefig(plot_filename)
print(f"预测结果图已保存到: {plot_filename}")

